#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库迁移工具
用于将user_id字段从手动配置改为基于IP地址的自动识别
"""

import logging
from datetime import datetime
from pymongo import MongoClient
from pymongo.errors import PyMongoError
from typing import Dict, List, Optional
from .ip_utils import generate_user_id_from_ip

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DatabaseMigration:
    """数据库迁移管理器"""
    
    def __init__(self, mongo_uri: str = "mongodb://localhost:27017/", db_name: str = "civil_code_db"):
        """
        初始化迁移管理器
        
        Args:
            mongo_uri: MongoDB连接URI
            db_name: 数据库名称
        """
        self.mongo_uri = mongo_uri
        self.db_name = db_name
        self.client = None
        self.db = None
        
    def connect(self) -> bool:
        """
        连接数据库
        
        Returns:
            bool: 连接是否成功
        """
        try:
            self.client = MongoClient(self.mongo_uri)
            self.db = self.client[self.db_name]
            # 测试连接
            self.client.admin.command('ping')
            logger.info(f"成功连接到数据库: {self.db_name}")
            return True
        except PyMongoError as e:
            logger.error(f"数据库连接失败: {e}")
            return False
    
    def close(self):
        """关闭数据库连接"""
        if self.client:
            self.client.close()
            logger.info("数据库连接已关闭")
    
    def backup_collections(self) -> bool:
        """
        备份相关集合
        
        Returns:
            bool: 备份是否成功
        """
        try:
            collections_to_backup = ['user_favorites', 'search_history']
            backup_suffix = datetime.now().strftime("_%Y%m%d_%H%M%S_backup")
            
            for collection_name in collections_to_backup:
                if collection_name in self.db.list_collection_names():
                    source_collection = self.db[collection_name]
                    backup_collection_name = f"{collection_name}{backup_suffix}"
                    
                    # 复制数据到备份集合
                    documents = list(source_collection.find())
                    if documents:
                        self.db[backup_collection_name].insert_many(documents)
                        logger.info(f"备份集合 {collection_name} -> {backup_collection_name}, 文档数: {len(documents)}")
                    else:
                        logger.info(f"集合 {collection_name} 为空，跳过备份")
                else:
                    logger.info(f"集合 {collection_name} 不存在，跳过备份")
            
            return True
        except PyMongoError as e:
            logger.error(f"备份失败: {e}")
            return False
    
    def migrate_user_favorites(self) -> Dict[str, int]:
        """
        迁移用户收藏集合
        
        为现有的user_id生成对应的IP哈希标识，并添加user_ip字段
        
        Returns:
            Dict[str, int]: 迁移统计信息
        """
        collection = self.db['user_favorites']
        stats = {'total': 0, 'migrated': 0, 'skipped': 0, 'errors': 0}
        
        try:
            # 获取所有文档
            documents = list(collection.find())
            stats['total'] = len(documents)
            
            logger.info(f"开始迁移 user_favorites 集合，共 {stats['total']} 条记录")
            
            for doc in documents:
                try:
                    current_user_id = doc.get('user_id', 'anonymous')
                    
                    # 如果已经是IP格式的user_id，跳过
                    if current_user_id.startswith('ip_'):
                        stats['skipped'] += 1
                        continue
                    
                    # 为anonymous用户生成默认IP标识
                    if current_user_id == 'anonymous':
                        # 使用127.0.0.1作为匿名用户的默认IP
                        user_ip = '127.0.0.1'
                        new_user_id = self._generate_user_id_hash(user_ip)
                    else:
                        # 对于其他用户ID，尝试解析为IP或生成哈希
                        user_ip = current_user_id if self._is_valid_ip(current_user_id) else '127.0.0.1'
                        new_user_id = self._generate_user_id_hash(user_ip)
                    
                    # 更新文档
                    update_result = collection.update_one(
                        {'_id': doc['_id']},
                        {
                            '$set': {
                                'user_id': new_user_id,
                                'user_ip': user_ip,
                                'migrated_at': datetime.now(),
                                'original_user_id': current_user_id
                            }
                        }
                    )
                    
                    if update_result.modified_count > 0:
                        stats['migrated'] += 1
                        logger.debug(f"迁移收藏记录: {current_user_id} -> {new_user_id}")
                    else:
                        stats['errors'] += 1
                        logger.warning(f"迁移失败: 文档ID {doc['_id']}")
                        
                except Exception as e:
                    stats['errors'] += 1
                    logger.error(f"迁移单个文档失败: {e}")
            
            logger.info(f"user_favorites 迁移完成: {stats}")
            return stats
            
        except PyMongoError as e:
            logger.error(f"迁移 user_favorites 失败: {e}")
            return stats
    
    def migrate_search_history(self) -> Dict[str, int]:
        """
        迁移搜索历史集合
        
        Returns:
            Dict[str, int]: 迁移统计信息
        """
        collection = self.db['search_history']
        stats = {'total': 0, 'migrated': 0, 'skipped': 0, 'errors': 0}
        
        try:
            # 获取所有文档
            documents = list(collection.find())
            stats['total'] = len(documents)
            
            logger.info(f"开始迁移 search_history 集合，共 {stats['total']} 条记录")
            
            for doc in documents:
                try:
                    current_user_id = doc.get('user_id', 'anonymous')
                    
                    # 如果已经是IP格式的user_id，跳过
                    if current_user_id.startswith('ip_'):
                        stats['skipped'] += 1
                        continue
                    
                    # 为anonymous用户生成默认IP标识
                    if current_user_id == 'anonymous':
                        user_ip = '127.0.0.1'
                        new_user_id = self._generate_user_id_hash(user_ip)
                    else:
                        user_ip = current_user_id if self._is_valid_ip(current_user_id) else '127.0.0.1'
                        new_user_id = self._generate_user_id_hash(user_ip)
                    
                    # 更新文档
                    update_result = collection.update_one(
                        {'_id': doc['_id']},
                        {
                            '$set': {
                                'user_id': new_user_id,
                                'user_ip': user_ip,
                                'migrated_at': datetime.now(),
                                'original_user_id': current_user_id
                            }
                        }
                    )
                    
                    if update_result.modified_count > 0:
                        stats['migrated'] += 1
                        logger.debug(f"迁移搜索历史: {current_user_id} -> {new_user_id}")
                    else:
                        stats['errors'] += 1
                        
                except Exception as e:
                    stats['errors'] += 1
                    logger.error(f"迁移单个文档失败: {e}")
            
            logger.info(f"search_history 迁移完成: {stats}")
            return stats
            
        except PyMongoError as e:
            logger.error(f"迁移 search_history 失败: {e}")
            return stats
    
    def create_indexes(self) -> bool:
        """
        创建新的索引结构
        
        Returns:
            bool: 索引创建是否成功
        """
        try:
            # user_favorites 索引
            favorites_collection = self.db['user_favorites']
            favorites_collection.create_index([("user_id", 1), ("article_id", 1)], unique=True, name="user_article_unique")
            favorites_collection.create_index([("user_ip", 1)], name="user_ip_index")
            favorites_collection.create_index([("created_at", -1)], name="created_at_index")
            
            # search_history 索引
            history_collection = self.db['search_history']
            history_collection.create_index([("user_id", 1), ("query", 1)], unique=True, name="user_query_unique")
            history_collection.create_index([("user_ip", 1)], name="user_ip_index")
            history_collection.create_index([("created_at", -1)], name="created_at_index")
            
            logger.info("数据库索引创建完成")
            return True
            
        except PyMongoError as e:
            logger.error(f"创建索引失败: {e}")
            return False
    
    def _is_valid_ip(self, ip_str: str) -> bool:
        """验证IP地址格式"""
        import ipaddress
        try:
            ipaddress.ip_address(ip_str)
            return True
        except ValueError:
            return False
    
    def _generate_user_id_hash(self, ip_address: str) -> str:
        """
        基于IP地址生成用户标识符（不依赖Flask上下文）
        
        Args:
            ip_address: IP地址
            
        Returns:
            str: 用户标识符
        """
        import hashlib
        hash_obj = hashlib.sha256(ip_address.encode('utf-8'))
        hash_hex = hash_obj.hexdigest()
        return f"ip_{hash_hex[:12]}"
    
    def get_migration_status(self) -> Dict[str, any]:
        """
        获取迁移状态
        
        Returns:
            Dict: 迁移状态信息
        """
        status = {
            'favorites': {'total': 0, 'migrated': 0, 'ip_format': 0},
            'history': {'total': 0, 'migrated': 0, 'ip_format': 0},
            'collections_exist': []
        }
        
        try:
            # 检查集合是否存在
            collections = self.db.list_collection_names()
            status['collections_exist'] = collections
            
            # 检查 user_favorites 状态
            if 'user_favorites' in collections:
                fav_collection = self.db['user_favorites']
                status['favorites']['total'] = fav_collection.count_documents({})
                status['favorites']['migrated'] = fav_collection.count_documents({'migrated_at': {'$exists': True}})
                status['favorites']['ip_format'] = fav_collection.count_documents({'user_id': {'$regex': '^ip_'}})
            
            # 检查 search_history 状态
            if 'search_history' in collections:
                hist_collection = self.db['search_history']
                status['history']['total'] = hist_collection.count_documents({})
                status['history']['migrated'] = hist_collection.count_documents({'migrated_at': {'$exists': True}})
                status['history']['ip_format'] = hist_collection.count_documents({'user_id': {'$regex': '^ip_'}})
            
            return status
            
        except PyMongoError as e:
            logger.error(f"获取迁移状态失败: {e}")
            return status
    
    def run_full_migration(self) -> bool:
        """
        执行完整迁移流程
        
        Returns:
            bool: 迁移是否成功
        """
        logger.info("开始执行完整数据库迁移...")
        
        # 连接数据库
        if not self.connect():
            return False
        
        try:
            # 1. 备份数据
            logger.info("步骤 1: 备份现有数据")
            if not self.backup_collections():
                logger.error("备份失败，终止迁移")
                return False
            
            # 2. 迁移用户收藏
            logger.info("步骤 2: 迁移用户收藏数据")
            fav_stats = self.migrate_user_favorites()
            
            # 3. 迁移搜索历史
            logger.info("步骤 3: 迁移搜索历史数据")
            hist_stats = self.migrate_search_history()
            
            # 4. 创建索引
            logger.info("步骤 4: 创建数据库索引")
            if not self.create_indexes():
                logger.warning("索引创建失败，但迁移数据仍然有效")
            
            # 5. 验证迁移结果
            logger.info("步骤 5: 验证迁移结果")
            status = self.get_migration_status()
            logger.info(f"迁移完成状态: {status}")
            
            logger.info("数据库迁移成功完成！")
            return True
            
        except Exception as e:
            logger.error(f"迁移过程中发生错误: {e}")
            return False
        finally:
            self.close()


def main():
    """主函数，执行迁移"""
    migration = DatabaseMigration()
    success = migration.run_full_migration()
    
    if success:
        print("✅ 数据库迁移成功完成")
    else:
        print("❌ 数据库迁移失败")
        exit(1)


if __name__ == "__main__":
    main()
